#include <assert.h>
#include <ctype.h>
#include <stdio.h>
#include <omp.h>
#include <mpi.h>
#include <stdlib.h>
#include <string.h>
#include "sieve.h"
#include "util.h"
#include "word_map.h"

#define MAX_PROCESSES 50

/* The root process' rank */
#define ROOT 0

/* Tags used in messages */
#define TAG_END 0
#define TAG_WORD 1

static int find_displacements(const char *buffer,
                              int buffer_size,
                              int processes,
                              int counts[],
                              int displacements[]);
static int preprocess_file(char *buffer, int *max_word_length, int *max_sum);
static char *process_file(const char *filename,
                          int processes,
                          int *max_word_length,
                          int *max_sieve,
                          int counts[],
                          int displacements[]);
static void process_words(char *words_buffer, long words_buffer_size);
static void process_words_root(int processes, int max_word_length);
static char *read_file(const char *filename);
static void run_non_root();
static void run_root(const char *filename);

/*
 * Finds out the correct displacements in buffer respecting the non-word split
 * along with the respective count (number of characters in the displacement).
 * Returns the largest displacement size.
 *
 * FIXME: If the number of processes is too high, the text is broken apart in
 * weird ways.
 */
static int
find_displacements(const char *buffer,
                   int buffer_size,
                   int processes,
                   int counts[],
                   int displacements[])
{
    int i, largest_count;

    largest_count = 0;
    displacements[0] = counts[0] = 0;
    displacements[processes] = buffer_size;
    /* Split into proccesses-1 parts since the root does not process a piece */
    for (i = processes - 1; i > 0; --i)
        displacements[i] = buffer_size * (i - 1) / (processes - 1);

#pragma omp parallel for
    for (i = 1; i < processes; ++i) {
        counts[i] = displacements[i + 1] - displacements[i];
        if (counts[i] > largest_count)
            largest_count = counts[i];
    }

    /* Do not split words; adjust displacements accordingly */
    for (i = 2; i < processes; ++i)
        while (buffer[displacements[i]] != ' ') {
            --displacements[i];
            ++counts[i];
            --counts[i - 1];
        }

    return largest_count;
}

/*
 * Strips non 0-9, A-Z or a-z characters from the given string. Also finds out
 * the maximum word length, the maximum value that is the sum of all the
 * characters in a word and returns the new buffer length.
 */
static int
preprocess_file(char *buffer, int *max_word_length, int *max_sum)
{
    int i, sum, word_length;
    char *ptr;

    i = 0;
    ptr = buffer;
    *max_sum = 0;
    *max_word_length = 0;
    while (buffer[i] != '\0') {
        /* Skip all non alnum characters */
        while (!isalnum(buffer[i]) && buffer[i] != '\0')
            ++i;

        sum = 0;
        word_length = 0;
        while (isalnum(buffer[i]) && buffer[i] != '\0') {
            *ptr++ = buffer[i];
            sum += buffer[i];
            ++word_length;
            ++i;
        }

        if (word_length > *max_word_length)
            *max_word_length = word_length;
        if (sum > *max_sum)
            *max_sum = sum;
        if (buffer[i] != '\0')
            *ptr++ = ' '; /* Not the end yet; add a space and keep going */
    }
    *ptr = '\0';

    return (int) (ptr - buffer);
}

/*
 * Prints all the words in the given map.
 */
static void
print_words(const WordMap *map)
{
    int sum;
    char *ptr;

    if (map != NULL) {
        print_words(map->left);
        sum = 0;
        for (ptr = map->str; *ptr; ptr++)
            sum += *ptr;
        printf("%s - %d - %s\n", map->str, sum, is_prime(sum) ? "sim" : "não");
        print_words(map->right);
    }
}

/*
 * Reads the input file, pre-processes it and returns the input file buffer.
 * The parameter outputs are:
 * - maximum word length;
 * - maximum number to run the sieve of Eratosthenes;
 * - the number of elements in each buffer displacement;
 * - the buffer displacements.
 */
static char *
process_file(const char *filename,
             int processes,
             int *max_word_length,
             int *max_sieve,
             int counts[],
             int displacements[])
{
    long buffer_size;
    char *buffer;

    assert(processes > 1 && processes <= MAX_PROCESSES);

    buffer = read_file(filename);
    assert(buffer != NULL);

    buffer_size = preprocess_file(buffer, max_word_length, max_sieve);
    find_displacements(buffer, buffer_size, processes, counts, displacements);
    return buffer;
}

/*
 * For each word in words_buffer, check whether it is a palindrome and then send
 * to root.
 */
static void
process_words(char *words_buffer, long words_buffer_size)
{
    int word_length;
    char *word;

    word = strtok(words_buffer, " ");
    while (word != NULL) {
        word_length = strlen(word);
        strtolower(word);
        if (word_length > 2 && is_palindrome(word, word_length))
            MPI_Send(word, word_length + 1, MPI_CHAR,
                     ROOT, TAG_WORD, MPI_COMM_WORLD);
        word = strtok(NULL, " ");
    } 

    /* We are done; notify the root process */
    MPI_Send(NULL, 0, MPI_CHAR, ROOT, TAG_END, MPI_COMM_WORLD);
}

/*
 * Add each received palindrome to a binary tree, ignoring repeated ones. When
 * all processes finish sending palindromes, print all palindromes ordered
 * lexicographically, along with some text indicating whether the sum of the
 * characters is a prime number.
 */
static void
process_words_root(int processes, int max_word_length)
{
    WordMap *map;
    int processes_alive;
    MPI_Status status;
    char recvbuffer[max_word_length + 1];
    int recvbuffer_size;

    recvbuffer_size = max_word_length + 1;
    processes_alive = processes - 1;
    map = word_map_new();
    do {
        MPI_Recv(recvbuffer, recvbuffer_size, MPI_CHAR,
                 MPI_ANY_SOURCE, MPI_ANY_TAG, MPI_COMM_WORLD, &status);
        switch (status.MPI_TAG) {
        case TAG_END:
            processes_alive--;
            break;
        case TAG_WORD:
            word_map_add(&map, recvbuffer);
            break;
        }
    } while (processes_alive > 0);

    print_words(map);
    word_map_free(map);
}

/*
 * Reads the file named filename into a buffer and returns that buffer's address.
 */
static char *
read_file(const char *filename)
{
    long file_size;
    FILE *file;
    char *buffer;

    buffer = NULL;
    file = fopen(filename, "r");
    if (file == NULL)
        goto read_file_error;

    /* Find out the file size */
    fseek(file, 0, SEEK_END);
    file_size = ftell(file);
    fseek(file, 0, SEEK_SET);

    /* Allocate a buffer to read the file into */
    buffer = (char *) malloc(sizeof(char) * file_size + 1);
    if (buffer == NULL)
        goto read_file_error;

    /* Read the file into the allocated buffer */
    fread(buffer, sizeof(char), file_size, file);
    if (ferror(file))
        goto read_file_error;

    if (buffer[file_size - 1] == '\n')
        buffer[file_size - 1] = '\0';
    else
        buffer[file_size] = '\0';
    fclose(file);
    return buffer;

read_file_error:
    if (buffer != NULL)
        free(buffer);
    fclose(file);
    return NULL;
}

static void
run_non_root()
{
    char *recvbuffer;
    int recvbuffer_size;

    /* Obtain the buffer size from broadcast */
    MPI_Bcast(&recvbuffer_size, 1, MPI_INT, ROOT, MPI_COMM_WORLD);

    recvbuffer = (char *) malloc(sizeof(char) * (recvbuffer_size + 1));
    assert(recvbuffer != NULL);

    /* Obtain a piece of the file scattered */
    MPI_Scatterv(NULL, NULL, NULL, MPI_CHAR,
        recvbuffer, recvbuffer_size, MPI_CHAR, ROOT, MPI_COMM_WORLD);

    process_words(recvbuffer, recvbuffer_size);
    free(recvbuffer);
}

static void
run_root(const char *filename)
{
    char *buffer;
    int max_count, max_sieve, max_word_length, processes,
        counts[MAX_PROCESSES], displacements[MAX_PROCESSES];

    MPI_Comm_size(MPI_COMM_WORLD, &processes);
    buffer = process_file(filename, processes, &max_word_length,
                 &max_sieve, counts, displacements);

    /* Obtain the maximum displacement size and... */
    max_count = maxv(counts, processes);
    /* ...broadcast it so that all processes can allocate enough memory */
    MPI_Bcast(&max_count, 1, MPI_INT, ROOT, MPI_COMM_WORLD);

    /* Scatter the file across all processes, skipping the root */
    MPI_Scatterv(buffer, counts, displacements, MPI_CHAR,
                 MPI_IN_PLACE, 0, MPI_CHAR, ROOT, MPI_COMM_WORLD);

    free(buffer);

    sieve_init(max_sieve);
    process_words_root(processes, max_word_length);
    sieve_finalize();
}

int
main(int argc, char *argv[])
{
    int own_rank;

    if (argc != 2)
        exit(EXIT_FAILURE);

    MPI_Init(&argc, &argv);
    MPI_Comm_rank(MPI_COMM_WORLD, &own_rank);
    if (own_rank == ROOT)
        run_root(argv[1]);
    else
        run_non_root();
    MPI_Finalize();

    return EXIT_SUCCESS;
}
